﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Collections;

namespace DiffusionMonteCarlo
{
    public class PathIntegral
    {
        Input inp;
        PotentialEnergy<double[,], double> V;  // [numberOfParticles,SpaceDimension]
        
        public PathIntegral(Input input, PotentialEnergy<double[,], double> potentialEnergy)
        {
            inp = input;
            V = potentialEnergy;
            Initialize();
        }

        double ET = 0.0;
        int[] nn;
        double[] Psi;
        List<double[,]> Replica = new List<double[,]>();
        Normaldev dev;
        Random ran;
        private void Initialize()
        {
            nn = new int[inp.SpaceDimension];
            int numberOfSize = 1;
            for (int i = 0; i < inp.SpaceDimension; i++)
            {
                nn[i] = (int)Math.Floor((inp.LimitsOfCoordinates[i, 1]- inp.LimitsOfCoordinates[i, 0])
                    / inp.ValueOfLengthInterval[i]);
                for (int j = 0; j < inp.NumberOfParticles; j++) numberOfSize *= nn[i];
            }
            Psi = new double[numberOfSize];
            for (int i = 0; i < inp.InitialNumberOfReplicas; i++)
            {
                double[,] x = new double[inp.NumberOfParticles, inp.SpaceDimension];
                for (int j = 0; j < inp.NumberOfParticles; j++)
                    for (int k = 0; k < inp.SpaceDimension; k++)
                        x[j, k] = inp.InitialPosition[j, k];     
                Replica.Add(x);
            }
            dev = new Normaldev(0.0, 1.0, inp.Seed);
            ran = new Random(inp.Seed + 123L);
        }

        private void Walk()
        {
            double st = Math.Sqrt(inp.ValueOfTimeStep);
            for (int i = 0; i < Replica.Count; i++)
                for (int j = 0; j < inp.NumberOfParticles; j++)
                    for (int k = 0; k < inp.SpaceDimension; k++)
                        Replica[i][j,k] = Replica[i][j,k] + st * dev.dev();                
        }

        private int old;
        private void Branch()
        {
            old = Replica.Count;
            int cc = Replica.Count;
            for (int i = 0; i < cc; i++)
            {
                int mn = Replication(Replica[i]);
                if (mn == 0) { Replica.RemoveAt(i); i--; cc--; }
                else if (mn == 2) 
                {
                    double[,] x = new double[inp.NumberOfParticles, inp.SpaceDimension];
                    for (int j = 0; j < inp.NumberOfParticles; j++)
                        for (int k = 0; k < inp.SpaceDimension; k++)
                            x[j, k] = Replica[i][j, k];   
                    Replica.Add(x);
                }
                else if (mn == 3)
                {
                    double[,] x = new double[inp.NumberOfParticles, inp.SpaceDimension];
                    for (int j = 0; j < inp.NumberOfParticles; j++)
                        for (int k = 0; k < inp.SpaceDimension; k++)
                            x[j, k] = Replica[i][j, k];
                    Replica.Add(x);
                    double[,] y = new double[inp.NumberOfParticles, inp.SpaceDimension];
                    for (int j = 0; j < inp.NumberOfParticles; j++)
                        for (int k = 0; k < inp.SpaceDimension; k++)
                            y[j, k] = Replica[i][j, k];
                    Replica.Add(y);
                }
                else { }
            }               
        }
        private int Replication(double[,] x)
        {
            int mn = (int) Math.Floor(WeightFunction(x)+ran.doub());
            return ( mn < 3 ? mn : 3);            
        }
    
        private double sum = 0.0;
        private double WeightFunction(double[,] x)
        {
            double vx = V(x);
            sum += vx;
            return Math.Exp(-(vx - ET)*inp.ValueOfTimeStep);
        }

        private void Count()
        {
            ET = sum / old
                + 0.1 * (1.0 - (double) Replica.Count / (double) inp.InitialNumberOfReplicas);
            sum = 0.0;
        }

        private void FindPsi()
        {
            foreach (double[,] y in Replica)
            {
                int[,] n = new int[inp.NumberOfParticles, inp.SpaceDimension];
                for (int j = 0; j < inp.NumberOfParticles; j++)
                    for (int k = 0; k < inp.SpaceDimension; k++)
                    {
                        n[j, k] = (int)Math.Floor((y[j, k] - inp.LimitsOfCoordinates[k, 0]) / inp.ValueOfLengthInterval[k]);
                        if (n[j, k] < 0) n[j, k] = 0;
                        else if (n[j, k] >= nn[k]) n[j, k] = nn[k] - 1;
                        else { }
                    }
                Psi[SingleIndex(n)] += 1;
            }
        }
        private int SingleIndex(int[,] n)
        {
            int t = n[0, 0];
            for (int k = 1; k < inp.SpaceDimension; k++)
            {
                t *= nn[k];
                t += n[0, k];
            }
            for (int j = 1; j < inp.NumberOfParticles; j++)
            {
                for (int k = 0; k < inp.SpaceDimension; k++)
                {
                    t *= nn[k];
                    t += n[j, k];
                }
            }
            return t;
        }
        
        public Output Compute()
        {
            Output psi = new Output();
            int numberOfRun = (int)(inp.RunTime / inp.ValueOfTimeStep);
            for (int i = 0; i < numberOfRun; i++)
            {
                Walk();
                Branch();
                Count();
            }
            double sum = 0.0;
            for (int i = 0; i < numberOfRun; i++)
            {
                Walk();
                Branch();
                Count();
                sum += ET;
                FindPsi();
            } 
            psi.GroundStateEnergy = sum / (double)numberOfRun;

            double volume = 1.0;
            for (int i = 0; i < inp.SpaceDimension; i++) volume *= inp.ValueOfLengthInterval[i];

            double sumpsi = 0.0;
            for (int i = 0; i < Psi.Length; i++) sumpsi += Psi[i] * Psi[i];
            sumpsi *= Math.Pow(volume,(double)inp.NumberOfParticles);
            for (int i = 0; i < Psi.Length; i++) Psi[i] /= Math.Sqrt(sumpsi);
            psi.GroundStateWaveFunction = Psi;

            return psi;
        }
    }


    class Random
    {
        ulong u, v, w;
        public Random(ulong j)
        {
            v = 4101842887655102017L;
            w = 1;
            u = j ^ v; int64();
            v = u; int64();
            w = v; int64();
        }
        public ulong int64()
        {
            u = u * 2862933555777941757L + 7046029254386353087L;
            v ^= v >> 17; v ^= v << 31; v ^= v >> 8;
            w = 4294957665U * (w & 0xffffffff) + (w >> 32);
            ulong x = u ^ (u << 21); x ^= x >> 35; x ^= x << 4;
            return (x + v) ^ w;
        }
        public double doub()
        {
            return 5.42101086242752217E-20 * int64();
        }
        public uint int32()
        {
            return (uint)int64();
        }
        ulong breg;
        int bc;
        public byte int8()
        {
            if ((bc--) != 0) return (byte)(breg >>= 8);
            breg = int64();
            bc = 7;
            return (byte)breg;
        }
    }

    class Normaldev : Random
    {
        double mu, sig;
        public Normaldev(double mmu, double ssig, ulong i)
            : base(i)
        {
            mu = mmu;
            sig = ssig;
        }

        public virtual double dev()
        {
            double u, v, x, y, q;
            do
            {
                u = doub();
                v = 1.7156 * (doub() - 0.5);
                x = u - 0.449871;
                y = Math.Abs(v) + 0.386595;
                q = x * x + y * (0.19600 * y - 0.25472 * x);
            } while (q > 0.27597
                && (q > 0.27846 || v * v > -4.0 * Math.Log(u) * u * u));
            return mu + sig * v / u;
        }
    }
}
